package com.gitee.feizns.explore.data_structure.tree.cus.rbt;

import com.gitee.feizns.explore.data_structure.tree.cus.Node;
import com.gitee.feizns.explore.data_structure.tree.cus.PutRet;

import java.util.Collection;
import java.util.LinkedList;
import java.util.Queue;
import java.util.Set;
import java.util.TreeSet;
import java.util.stream.Collectors;

/**
 * @author feizns
 * @since 2019/11/27
 */
public class RbtImpl<K extends Comparable<K>, V> implements Rbt<RbtImpl.RBNode, K, V> {

    public static void main(String[] args) {
        RbtImpl<Integer, String> rbt = new RbtImpl<>();
        rbt.put(1, "1");
        rbt.put(2, "1");
        rbt.put(3, "1");
        rbt.put(4, "1");
        rbt.put(5, "1");
        rbt.put(6, "1");
//        rbt.put(7, "1");
        System.out.println(rbt);
    }

    private RBNode<K, V> root;

    private int size;

    @Override
    public Node root() {
        return root;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public V get(Object key) {
        RBNode<K, V> node = find(root, (K) key);
        return node != null ? node.val() : null;
    }

    @Override
    public V put(K key, V value) {
        PutRet<RBNode, V> ret = put(root, key, value);
        root = (RBNode<K, V>) ret.node();
        root.color = BLACK;
        return ret.v();
    }

    private PutRet<RBNode, V> put(RBNode<K, V> root, K k, V v) {
        if ( root == null ) {
            size++;
            return new PutRet<>(new RBNode(k, v), null);
        }

        int comp = k.compareTo(root.k);

        V old;
        if ( comp < 0 ) {
            PutRet<RBNode, V> ret = put(root.left, k, v);
            root.left = ret.node();
            old = ret.v();
        } else if ( comp > 0 ) {
            PutRet<RBNode, V> ret = put(root.right, k, v);
            root.right = ret.node();
            old = ret.v();
        } else {
            old = root.val();
            root.val(v);
        }

        //左黑、右红
        if ( !isR(root.left) && isR(root.right) )
            root = leftRotate(root);

        //左红、左左红
        if ( isR(root.left) && isR(root.left.left) )
            root = rightRotate(root);

        //左红、右红
        if ( isR(root.left) && isR(root.right) )
            flipColors(root);

        return new PutRet<>(root, old);
    }

    private static RBNode leftRotate(RBNode node) {
        RBNode right = node.right;
        node.right = right.left;
        right.left = node;
        right.color = node.color;
        node.color = RED;
        return right;
    }

    private static RBNode rightRotate(RBNode node) {
        RBNode left = node.left;
        node.left = left.right;
        left.right = node;
        left.color = node.color;
        node.color = RED;
        return left;
    }

    private static RBNode flipColors(RBNode node) {
        node.color = RED;
        node.left.color = BLACK;
        node.right.color = BLACK;
        return node;
    }

    private static boolean isR(RBNode node) {
        return node != null && node.color == RED;
    }

    private RBNode<K, V> find(RBNode<K, V> node, K k) {
        if ( node != null ) {
            int comp = k.compareTo(node.k);
            if ( comp < 0 && node.left != null )
                return find(node.left, k);
            else if ( comp > 0 && node.right != null )
                return find(node.right, k);
        }
        return node;
    }

    @Override
    public V remove(Object key) {
        return null;
    }

    @Override
    public void clear() {
        this.root = null;
        this.size = 0;
    }

    @Override
    public Set<K> keySet() {
        return entrySet().stream().map(Entry::getKey).collect(Collectors.toSet());
    }

    @Override
    public Collection<V> values() {
        return entrySet().stream().map(Entry::getValue).collect(Collectors.toList());
    }

    @Override
    public Set<Entry<K, V>> entrySet() {
        Set<Entry<K, V>> entries = new TreeSet<>();
        recursiveMiddleEach((k, v) -> entries.add(new EntryImpl(k, v)));
        return entries;
    }

    @Override
    public String toString() {
        Queue<RBNode> queue = new LinkedList<>();
        if ( root != null )
            queue.add(root);
        return toString(queue, new StringBuilder());
    }

    private static String toString(Queue<RBNode> queue, StringBuilder ret) {
        Queue<RBNode> newQueue = new LinkedList<>();
        while ( !queue.isEmpty() ) {
            RBNode node = queue.remove();
            ret.append(String.format("%s|%s ", node.k, isR(node) ? "R" : "B"));
            if ( node.left != null )
                newQueue.add(node.left);
            if ( node.right != null )
                newQueue.add(node.right);
        }
        ret.append("\n");
        if ( !newQueue.isEmpty() )
            toString(newQueue, ret);
        return ret.toString();
    }

    private static boolean BLACK = true;

    private static boolean RED = false;

    static class RBNode<K extends Comparable<K>, V> implements Node<RBNode<K, V>, K, V> {

        K k;

        V v;

        boolean color;

        RBNode<K, V> left, right;

        public RBNode(K k, V v) {
            this.k = k;
            this.v = v;
            this.color = RED;
        }

        @Override
        public K key() {
            return k;
        }

        @Override
        public void key(K k) {
            this.k = k;
        }

        @Override
        public V val() {
            return v;
        }

        @Override
        public void val(V v) {
            this.v = v;
        }

        @Override
        public RBNode<K, V> left() {
            return left;
        }

        @Override
        public void left(RBNode<K, V> left) {
            this.left = left;
        }

        @Override
        public RBNode<K, V> right() {
            return right;
        }

        @Override
        public void right(RBNode<K, V> right) {
            this.right = right;
        }

    }

    public static class EntryImpl<K, V> implements Entry<K, V> {

        K k;

        V v;

        public EntryImpl(K k, V v) {
            this.k = k;
            this.v = v;
        }

        @Override
        public K getKey() {
            return k;
        }

        @Override
        public V getValue() {
            return v;
        }

        @Override
        public V setValue(V value) {
            V old = v;
            this.v = value;
            return old;
        }

    }

}
